import numpy as np
import gym
import time
from nes_py.wrappers import JoypadSpace
import gym_super_mario_bros
from gym.spaces import Box
from gym import Wrapper
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT, COMPLEX_MOVEMENT, RIGHT_ONLY
import pygame

ACTIONS = [
	['NOOP'],
	['left'],
	['left', 'A'],
	['right'],
	['right', 'A'],
	['A'],
]

env = gym_super_mario_bros.make("SuperMarioBros-1-2-v0")
env = JoypadSpace(env, ACTIONS)
env.reset()

print(env.ram.shape)

'''
pygame.init()
pygame.display.set_mode((256, 240))

while True:
	env.render()
	time.sleep(1./60.)
	key_pressed = pygame.key.get_pressed()
	action = 0
	if key_pressed[pygame.K_a]:
		if key_pressed[pygame.K_SPACE]: 
			action = 2 
		else: 
			action = 1
	elif key_pressed[pygame.K_d]:
		if key_pressed[pygame.K_SPACE]: 
			action = 4 
		else: 
			action = 3
	elif key_pressed[pygame.K_SPACE]:
		action = 5
	if key_pressed[pygame.K_ESCAPE]:
		pygame.quit()
		break
	state, reward, done, info = env.step(action)
	if done:
		env.reset()
'''